%% On the Optimal Design of Transfers and Income-Tax Progressivity
% Numerical solution of the simple model
% This file as it is produces Figure 1 in the paper
% To find a more accurate solution for the benchmark case, change the Tgrid as described when it is generated

clear; clc; close all;
addpath('../Matlab_Figure_Files');

%% Parameters and Grids 

% Parameters
phi = 1/0.4;            % inverse Frisch (standard value)
tau = 0.181;            % benchmark progressivity (from HSV 2017)
n_0_target = 0.3;       % target for labor supply without transfers
B = n_0_target^(-1-phi)*(1-tau);    % labor disutility parameter to match target for labor supply
G = 0.233*n_0_target;   % spending (targeted spending output ratio 0.233)
rhoa = 0.935;           % persistence idiosyncratic shock
v_log_con = 0.18;                           % variance log consumption data (from HSV 2017)
vw   = (1-rhoa^2)*v_log_con*(1/(1-tau)^2);  % variance idiosyncratic shock to match variance of log consumption in the data
vw_aux = vw/(1-rhoa^2);                     % translation from AR(1) to lognormal shocks
disp('Calibration')
disp(['Implied vw to match variance of log consumption is = ',num2str(vw)])
disp(['Expressed for lognormal shocks rather than AR(1) = ',num2str(vw_aux)])
disp(' ')

% Tax grid
% for finding optimum with higher precision: finer grid from -0.10 to -0.08 with 111 points
tau_min = -0.35;                % minimum for tax progressivity grid
tau_max = 0.60;                 % maximum for tax progressivity grid
tau_n = 1000;                   % number of points in tax progressivity grid
tau_grid = linspace(tau_min,tau_max,tau_n)';    % tax progressivity grid

% Transfer grid
% for finding optimum with higher precision: finer grid from 0.065 to 0.075 with 50 points
T_min = -0.10;          % minimum for transfer grid
T_max = 0.10;           % maximum for transfer grid
T_n = 41;               % number of points in transfer grid
T_grid = linspace(T_min,T_max,T_n)';            % transfer grid

% Stack all combinations of taxes and transfers
fiscal = gridmake(tau_grid,T_grid);             % column 1: tax progressivity; column 2: transfer

% Uninsurable shocks
nodes_alpha_n = 15;                 % number of points for idiosyncratic shock
alpha_mean = -vw/(2*(1-rhoa^2));    % mean of idiosyncratic shock process (in logs)
alpha_var = vw/(1-rhoa^2);          % variance of idiosyncratic shock process
[nodes_alpha,weights_alpha] = qnwnorm(nodes_alpha_n,alpha_mean,alpha_var);  % grid and mass per grid point for idiosyncratic shock

%% Labor Supply, Consumption, Government Budget

tol = 1e-11;                                % tolerance level (common for bisection to clear govt budget and to find labor supply of HHs)
max_iter = 5000;                            % maximum number of iterations (common for bisection to clear govt budget and to find labor supply of HHs)
lambda_grid = zeros(T_n*tau_n,1);           % vector to store lambda for each tax system characterized by (tau,T)
ymean = zeros(T_n*tau_n,1);                 % vector to store mean income for each tax system characterized by (tau,T)
c_mat = zeros(nodes_alpha_n,T_n*tau_n);     % matrix to store consumption for each productivity level alpha and each tax system characterized by (tau,T)
n_mat = zeros(nodes_alpha_n,T_n*tau_n);     % matrix to store labor for each productivity level alpha and each tax system characterized by (tau,T)
y_mat = zeros(nodes_alpha_n,T_n*tau_n);     % matrix to store income for each productivity level alpha and each tax system characterized by (tau,T)
yhat_mat = zeros(nodes_alpha_n,T_n*tau_n);  % matrix to store auxiliary yhat for each productivity level alpha and each tax system characterized by (tau,T)
tax_mat = zeros(nodes_alpha_n,T_n*tau_n);   % matrix to store tax paid for each productivity level alpha and each tax system characterized by (tau,T)

for i = 1:T_n*tau_n                         % loop over all tax systems
        
    lambda_min = 0;                         % minimum for tax level parameter lambda
    lambda_max = 2;                         % maximum for tax level parameter lambda
    for iter_fiscal = 1:max_iter            % bisection to find lambda
        if iter_fiscal == max_iter          % report if max number of iterations is reached for finding lambda
            disp('Maximum number of iterations reached for lambda')
            disp(['Fiscal system indexed = ', num2str(i)])
            disp(['gamma = ',num2str(fiscal(i,1)),'T = ',num2str(fiscal(i,2))])
            disp(['error = ',num2str(err_fiscal)])
        end 
        
        lambda = 0.5*(lambda_min+lambda_max);   % update guess for lambda
        
        % This condition takes care of cases in which T and tau are large, which can be inconsistent with govt budget clearing
        if lambda < 0.1
            n_mat(:,i) = 50;    % set labor supply very high: will not choose this system
            c_mat(:,i) = 1e-8;  % set consumption very low: will not choose this system
            break
        end
        
        % Given taxes, solve for labor supply at all productivity levels
        for j = 1:nodes_alpha_n
            
            n_min = 0;  % minimum for bisection on labor supply
            n_max = 50; % maximum for bisection on labor supply
            for iter_alpha = 1:max_iter
                if iter_alpha == max_iter   % report if max number of iterations is reached when trying to find labor supply
                    disp('Maximum number of iterations reached for finding n as a function of alpha')
                    disp(['Fiscal system indexed ', num2str(i)])
                    disp(['alpha indexed ', num2str(j)])
                end
                    
                n = 0.5*(n_min+n_max);  % update guess for labor supply
                err_n = B*n^(1+phi) + (fiscal(i,2)/lambda)*B*n^(phi+fiscal(i,1))*exp(-(1-fiscal(i,1))*nodes_alpha(j)) - (1-fiscal(i,1));    % labor FOC
                if abs(err_n) < tol     % FOC is solved
                    n_mat(j,i) = n;     % store labor supply
                    c_mat(j,i) = lambda*(exp(nodes_alpha(j))*n)^(1-fiscal(i,1)) + fiscal(i,2);  % store consumption
                    break
                elseif err_n < 0        % update bounds for bisection on labor supply
                    n_min = n;
                else
                    n_max = n;
                end
            end
            
            % Store more variables
            y_mat(j,i) = n_mat(j,i)*exp(nodes_alpha(j));        % store income
            yhat_mat(j,i) = lambda*y_mat(j,i)^(1-fiscal(i,1));  % store auxiliary yhat 
            tax_mat(j,i) = y_mat(j,i)-yhat_mat(j,i);            % store tax payment
                        
        end
        
        Y = sum(weights_alpha.*y_mat(:,i));             % compute average income
        Yhat = sum(weights_alpha.*yhat_mat(:,i));       % compute average yhat
        ymean(i) = Y;                                   % store mean income
        
        err_fiscal = G + fiscal(i,2) - (Y-Yhat);        % check government budget constraint
    
        if abs(err_fiscal) < tol            % govt budget is cleared
            lambda_grid(i) = lambda;        % store lambda       
            break
        elseif err_fiscal < 0               % update bounds for lambda
            lambda_min = lambda;
        else
            lambda_max = lambda;
        end

    % End bisection to find lambda       
    end
% End loop over tax and transfer systems
end

% Check that consumption and labor supply are nonnegative
if min(min(c_mat)) < 0
    disp('Consumption negative')
end
if min(min(n_mat)) < 0
    disp('Labor negative')
end

%% Welfare

% Compute individual utilities and welfare
util_mat = log(c_mat) - (B/(1+phi))*n_mat.^(1+phi);     % matrix of individual utilities for each productivity level alpha and each tax system characterized by (tau,T)     
W = zeros(T_n*tau_n,1);
for i = 1:T_n*tau_n
    W(i) = sum(weights_alpha.*util_mat(:,i));           % vector of utilitarian welfare associated with each tax and transfer system
end

% Find optimal tax and transfer system
[~,loc] = max(W);       % position of tax and transfer system associated with highest welfare
disp('Optimal progressivity and lump sum')
disp(['Optimal progressivity is :', num2str(fiscal(loc,1))])
disp(['Optimal lambda is :', num2str(lambda_grid(loc))])
disp(['Optimal lump sum is :', num2str(fiscal(loc,2))])
disp(['Optimal lump sum relative to mean income is :', num2str(fiscal(loc,2)/ymean(loc))])
disp(['Optimal lump sum relative to calibrated mean income is :', num2str(fiscal(loc,2)/n_0_target)])
disp(['ymean under this system is :', num2str(ymean(loc))])
disp(' ')

% Optimal tau by T
opt_tau_vec = zeros(T_n,1);     % vector for the optimal tau by T
opt_welf_vec = zeros(T_n,1);    % vector for welfare associated with optimal tau by T
opt_ymean_vec = zeros(T_n,1);   % vector for mean income associated with optimal tau by T
opt_lambda = zeros(T_n,1);      % vectore for lambda associated with optimal tau by T
for i = 1:T_n                   % loop over all transfer levels
    W_aux = W((i-1)*tau_n+1:i*tau_n);                   % pick all welfare levels associated with different tau for current T
    fiscal_aux = fiscal((i-1)*tau_n+1:i*tau_n,:);       % pick all fiscal systems associated with different tau for current T
    ymean_aux = ymean((i-1)*tau_n+1:i*tau_n);           % pick all mean income levels associated with different tau for current T
    lambda_aux = lambda_grid((i-1)*tau_n+1:i*tau_n);    % pick all lambdas associated with different tau for current T 
    [~,loc] = max(W_aux);                               % find the location of the highest welfare for the current T
    opt_tau_vec(i) = fiscal_aux(loc,1);                 % pick the tau associated with the highest welfare for the current T
    opt_welf_vec(i) = W_aux(loc);                       % store the highest welfare for the current T
    opt_ymean_vec(i) = ymean_aux(loc);                  % pick the mean income associated with the highest welfare for the current T
    opt_lambda(i) = lambda_aux(loc);                    % pick the lambda associated with the highest welfare for the current T
    if abs(fiscal_aux(1,2)) < 1e-6                      % compute the optimal progressivity for the zero transfer case
        disp('Optimal progressivity')
        disp(['Optimal progressivity is :', num2str(fiscal_aux(loc,1))])
        disp(['Optimal lambda is :', num2str(lambda_aux(loc))])
        disp(['ymean under this system is :', num2str(ymean_aux(loc))])
        disp(' ')
        prog_noT = fiscal_aux(:,1);
        W_noT = W_aux;
    end
end

%% Optimal tau given T using formula

% Preallocations
W_formula = zeros(T_n,1);       % welfare from formula for each T given optimal choice of tau
tau_formula = zeros(T_n,1);     % optimal tau for each T implied by formula

% Variables for formula
tau_vec = tau_grid;                     % grid for tau to be evaluated with formula
n_0 = ((1-tau_vec)./B).^(1/(1+phi));    % labor supply without transfers given progressivity
vc = ((1-tau_vec).*vw)./(2*(1-rhoa^2)); % auxiliary variable
eff_nt = log(n_0-G);                    % efficiency term without transfers
ls_nt = -(1-tau_vec)./(1+phi);          % labor disutility term without transfers
rd_nt = -(1-tau_vec).*vc;               % redistribution term without transfers
W_nt_ha = eff_nt + ls_nt + rd_nt;       % welfare without transfers

for i = 1:T_n                           % loop over all transfer levels       
    T = T_grid(i);                      % pick current T
    Omega_e_ra = -(n_0./(n_0-G)).*(1/(1+phi)).*(1./(n_0-G)) + ((1-tau_vec)./(1+phi)).*(1./(n_0-G)); % efficiency term representative agent
    Omega_e_ha = ((tau_vec.*(1-tau_vec))./(n_0-G)).*(vw/(1-rhoa^2)).*(n_0./(n_0-G)).*(1/(1+phi));   % efficiency term heterogeneous agent
    Omega_r = (((1-tau_vec).^2)./(n_0-G)).*(vw/(1-rhoa^2));                                         % redistribution term
    W_wt_ha = W_nt_ha + T*(Omega_e_ra+Omega_e_ha+Omega_r);                                          % welfare with transfers
    [W_formula(i),aux_ha_wt] = max(W_wt_ha);        % pick best welfare level implied by formula for this T
    tau_formula(i) = tau_vec(aux_ha_wt);            % pick tau associated with this best welfare level for this T
    if abs(T) < 1e-6                                % also store what formula implies for zero transfer
        prog_noT_formula = tau_vec;                 % optimal progressivity implied by formula for zero transfer
        W_noT_formula = W_wt_ha;                    % welfare from formula with optimal progressivity for zero transfer
    end
end


%% Plots for paper

Tymean_grid = T_grid./opt_ymean_vec;    % transfer relative to mean income (where mean income is that of the respective economy)   
Tymean_calib = T_grid./n_0_target;      % transfer relative to mean income (where mean income is that of the calibrated economy)

fig = figure(101);
yyaxis left
plot(Tymean_calib*100,opt_tau_vec,'--','Color',[0    0.4470    0.7410],'Linewidth',2); hold on
plot(Tymean_calib*100,tau_formula,':','Color',[0    0.4470    0.7410],'Linewidth',2); hold off
set(gca,'XGrid','off','YGrid','on','Fontsize',12)  
set(gca,'TickLabelInterpreter','LaTex')
xlabel('$T$ in percent of calibrated mean income','Interpreter','LaTex','Fontsize',12)
xlim([min(Tymean_calib*100) max(Tymean_calib*100)])
ylabel('$\tau$','Interpreter','LaTex','Fontsize',12)
yyaxis right
plot(Tymean_calib*100,opt_welf_vec-max(opt_welf_vec),'-','Color',[0.8500    0.3250    0.0980],'Linewidth',2);
ylabel('Normalized welfare','Interpreter','LaTex','Fontsize',12)
leg = legend('Optimal progressivity: global solution','Optimal progressivity: formula','Welfare: global solution','Location','South');
set(leg,'Interpreter','LaTex','Fontsize',12);
legend boxoff
run graph_extended.m;
saveas(gcf,'am_global.pdf')

fig = figure(102);
yyaxis left
plot(Tymean_calib*100,opt_tau_vec,'--','Color',[0 0 0],'Linewidth',2); hold on
plot(Tymean_calib*100,tau_formula,':','Color',[0 0 0],'Linewidth',2); hold off
set(gca,'XGrid','off','YGrid','on','Fontsize',12)  
set(gca,'TickLabelInterpreter','LaTex')
xlabel('$T$ in percent of calibrated mean income','Interpreter','LaTex','Fontsize',12)
xlim([min(Tymean_calib*100) max(Tymean_calib*100)])
ylabel('$\tau$','Interpreter','LaTex','Fontsize',12)
yyaxis right
plot(Tymean_calib*100,opt_welf_vec-max(opt_welf_vec),'-','Color',[0.5 0.5 0.5],'Linewidth',2);
ylabel('Normalized welfare','Interpreter','LaTex','Fontsize',12)
leg = legend('Optimal progressivity: global solution','Optimal progressivity: formula','Welfare: global solution','Location','South');
set(leg,'Interpreter','LaTex','Fontsize',12);
legend boxoff
ax = gca;
ax.YAxis(1).Color = '[0 0 0]';
ax.YAxis(2).Color = '[0.4 0.4 0.4]';
run graph_extended.m;
saveas(gcf,'am_global_bw.pdf')
saveas(fig,'am_global_bw','epsc')
